import os
import json
import time
from typing import List, Dict, Optional
from datetime import datetime

class KlineCache:
    def __init__(self, cache_dir: str = "cache"):
        """初始化K线缓存管理器
        Args:
            cache_dir: 缓存文件存储目录
        """
        self.cache_dir = cache_dir
        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
    
    def _get_cache_file(self, symbol: str, timeframe: int) -> str:
        """获取缓存文件路径
        Args:
            symbol: 交易对名称
            timeframe: K线周期（分钟）
        """
        return os.path.join(self.cache_dir, f"{symbol}_{timeframe}m.json")
    
    def load_cache(self, symbol: str, timeframe: int) -> List[Dict]:
        """从缓存文件加载K线数据
        Args:
            symbol: 交易对名称
            timeframe: K线周期（分钟）
        Returns:
            List[Dict]: K线数据列表
        """
        cache_file = self._get_cache_file(symbol, timeframe)
        if not os.path.exists(cache_file):
            return []
        
        try:
            with open(cache_file, 'r') as f:
                data = json.load(f)
                return data.get('klines', [])
        except Exception as e:
            print(f"加载缓存文件失败: {str(e)}")
            return []
    
    def save_cache(self, symbol: str, timeframe: int, klines: List[Dict]) -> None:
        """保存K线数据到缓存文件
        Args:
            symbol: 交易对名称
            timeframe: K线周期（分钟）
            klines: K线数据列表
        """
        cache_file = self._get_cache_file(symbol, timeframe)
        try:
            data = {
                'symbol': symbol,
                'timeframe': timeframe,
                'last_update': int(time.time()),
                'klines': klines
            }
            with open(cache_file, 'w') as f:
                json.dump(data, f, indent=2)
        except Exception as e:
            print(f"保存缓存文件失败: {str(e)}")
    
    def merge_klines(self, cached_klines: List[Dict], new_klines: List[Dict]) -> List[Dict]:
        """合并缓存的K线数据和新的K线数据
        Args:
            cached_klines: 缓存的K线数据
            new_klines: 新的K线数据
        Returns:
            List[Dict]: 合并后的K线数据
        """
        # 转换为字典格式，使用时间戳作为键
        kline_dict = {k['time']: k for k in cached_klines}
        
        # 合并新数据
        for kline in new_klines:
            kline_dict[kline['time']] = kline
        
        # 转换回列表并按时间排序
        merged = list(kline_dict.values())
        merged.sort(key=lambda x: x['time'])
        return merged
    
    def get_last_kline_time(self, symbol: str, timeframe: int) -> Optional[int]:
        """获取最后一根K线的时间戳
        Args:
            symbol: 交易对名称
            timeframe: K线周期（分钟）
        Returns:
            Optional[int]: 最后一根K线的时间戳，如果没有缓存数据则返回None
        """
        klines = self.load_cache(symbol, timeframe)
        if not klines:
            return None
        return max(k['time'] for k in klines)
    
    def clean_old_klines(self, klines: List[Dict], max_count: int = 1000) -> List[Dict]:
        """清理旧的K线数据，只保留最新的max_count根K线
        Args:
            klines: K线数据列表
            max_count: 最大保留的K线数量
        Returns:
            List[Dict]: 清理后的K线数据列表
        """
        if len(klines) <= max_count:
            return klines
        return sorted(klines, key=lambda x: x['time'], reverse=True)[:max_count] 